#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
改进的键盘检测模块 - 解决旋转和比例问题
"""

import cv2
import numpy as np
import numpy.typing as npt
from typing import Tuple, List, Dict, Optional, Any
import logging
from dataclasses import dataclass
from enum import Enum
import math
from enhanced_rotation_correction import EnhancedRotationCorrection

class KeyIssueType(Enum):
    MISSING = "缺失"
    MISPLACED = "位置错误"
    ROTATED = "方向错误"
    WRONG_KEY = "错误按键"

@dataclass
class KeyIssue:
    issue_type: KeyIssueType
    position: Tuple[int, int]
    description: str
    confidence: float

class ImprovedKeyboardDetector:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
        # 键盘检测参数 - 基于比例而非像素
        self.min_keyboard_area_ratio = 0.1  # 最小键盘区域面积比例
        self.key_contour_area_range_ratio = (0.0001, 0.01)  # 按键轮廓面积比例范围
        self.aspect_ratio_tolerance = 0.3  # 长宽比容差
        
        # 模板匹配参数
        self.match_threshold = 0.7  # 匹配阈值
        self.rotation_angles = [0, 90, 180, 270]  # 检测旋转角度
        
        # 比例检测参数
        self.position_tolerance_ratio = 0.05  # 位置容差比例（相对于图像尺寸）
        self.size_tolerance_ratio = 0.2  # 大小容差比例
        self.angle_tolerance = 10.0  # 角度容差（度）
        
        # 增强旋转校正
        self.rotation_corrector = EnhancedRotationCorrection()
        
    def detect_keyboard_keys(self, image: npt.NDArray) -> List[Dict[str, Any]]:
        """
        检测键盘上的所有按键 - 改进版本
        
        Args:
            image: 输入图像
            
        Returns:
            按键信息列表，每个按键包含位置、大小、角度等信息
        """
        # 1. 首先进行图像旋转校正
        corrected_image, rotation_angle = self._correct_image_rotation(image)
        
        # 2. 预处理图像
        processed = self._preprocess_for_key_detection(corrected_image)
        
        # 3. 检测按键轮廓
        key_contours = self._detect_key_contours(processed)
        
        # 4. 提取按键信息
        keys = []
        for contour in key_contours:
            key_info = self._extract_key_info(contour, corrected_image)
            if key_info:
                keys.append(key_info)
                
        # 5. 按位置排序（从上到下，从左到右）
        keys.sort(key=lambda k: (k['y'], k['x']))
        
        # 6. 计算键盘布局比例
        if keys:
            self._calculate_keyboard_layout_ratios(keys, corrected_image.shape)
        
        return keys, corrected_image, rotation_angle
    
    def auto_correct_reference_image(self, image: npt.NDArray) -> Tuple[npt.NDArray, float, dict]:
        """
        自动校正参考图像的旋转角度 - 专门用于参考图像摆正
        
        Args:
            image: 输入参考图像
            
        Returns:
            (校正后的图像, 旋转角度, 检测信息)
        """
        self.logger.info("开始自动校正参考图像...")
        
        # 使用增强的旋转校正
        corrected_image, rotation_angle, correction_info = self.rotation_corrector.auto_correct_reference_image(image)
        
        # 记录校正信息
        self.logger.info(f"参考图像校正完成:")
        self.logger.info(f"  - 旋转角度: {rotation_angle:.1f}°")
        self.logger.info(f"  - 校正方法: {correction_info.get('method', 'unknown')}")
        self.logger.info(f"  - 置信度: {correction_info.get('confidence', 0):.2f}")
        
        return corrected_image, rotation_angle, correction_info
        
    def _correct_image_rotation(self, image: npt.NDArray) -> Tuple[npt.NDArray, float]:
        """
        校正图像旋转 - 基于键盘边缘检测
        """
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        # 边缘检测
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
        
        # 使用霍夫变换检测直线
        lines = cv2.HoughLines(edges, 1, np.pi/180, threshold=100)
        
        if lines is not None and len(lines) > 0:
            # 计算主要方向的角度
            angles = []
            for line in lines:
                rho, theta = line[0]
                angle = theta * 180 / np.pi
                # 将角度标准化到 [-45, 45] 范围
                if angle > 45:
                    angle -= 90
                elif angle < -45:
                    angle += 90
                angles.append(angle)
            
            # 计算平均角度
            if angles:
                avg_angle = np.median(angles)
                
                # 如果角度超过阈值，进行旋转校正
                if abs(avg_angle) > 2:  # 2度阈值
                    # 旋转图像
                    h, w = image.shape[:2]
                    center = (w // 2, h // 2)
                    rotation_matrix = cv2.getRotationMatrix2D(center, -avg_angle, 1.0)
                    corrected_image = cv2.warpAffine(image, rotation_matrix, (w, h), 
                                                   flags=cv2.INTER_CUBIC, 
                                                   borderMode=cv2.BORDER_REPLICATE)
                    return corrected_image, avg_angle
        
        return image, 0.0
        
    def _preprocess_for_key_detection(self, image: npt.NDArray) -> npt.NDArray:
        """为按键检测预处理图像"""
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        # 自适应直方图均衡化
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)
        
        # 高斯模糊
        blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
        
        # 自适应阈值处理
        binary = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, 11, 2
        )
        
        # 形态学操作
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        
        return cleaned
        
    def _detect_key_contours(self, binary_image: npt.NDArray) -> List[npt.NDArray]:
        """检测按键轮廓 - 基于比例"""
        contours, _ = cv2.findContours(
            binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        
        # 计算图像总面积
        total_area = binary_image.shape[0] * binary_image.shape[1]
        
        # 过滤轮廓 - 基于面积比例
        key_contours = []
        for contour in contours:
            area = cv2.contourArea(contour)
            area_ratio = area / total_area
            
            if self.key_contour_area_range_ratio[0] <= area_ratio <= self.key_contour_area_range_ratio[1]:
                # 检查长宽比
                x, y, w, h = cv2.boundingRect(contour)
                aspect_ratio = w / h if h > 0 else 0
                
                if 0.5 <= aspect_ratio <= 2.0:  # 合理的长宽比范围
                    key_contours.append(contour)
                    
        return key_contours
        
    def _extract_key_info(self, contour: npt.NDArray, original_image: npt.NDArray) -> Optional[Dict[str, Any]]:
        """提取按键信息"""
        # 获取边界框
        x, y, w, h = cv2.boundingRect(contour)
        
        # 计算最小外接矩形
        rect = cv2.minAreaRect(contour)
        center, size, angle = rect
        
        # 提取按键区域
        key_region = original_image[y:y+h, x:x+w]
        
        # 计算按键特征
        key_info = {
            'x': int(center[0]),
            'y': int(center[1]),
            'width': int(size[0]),
            'height': int(size[1]),
            'angle': angle,
            'area': cv2.contourArea(contour),
            'region': key_region,
            'contour': contour
        }
        
        return key_info
        
    def _calculate_keyboard_layout_ratios(self, keys: List[Dict], image_shape: Tuple[int, int, int]):
        """计算键盘布局比例"""
        if not keys:
            return
            
        # 计算图像尺寸
        img_height, img_width = image_shape[:2]
        
        # 计算按键的平均尺寸
        avg_width = np.mean([key['width'] for key in keys])
        avg_height = np.mean([key['height'] for key in keys])
        
        # 计算按键间距
        if len(keys) > 1:
            # 按行分组
            rows = self._group_keys_by_rows(keys)
            if rows:
                # 计算行间距
                row_spacing = self._calculate_row_spacing(rows)
                # 计算列间距
                col_spacing = self._calculate_col_spacing(rows)
                
                # 存储比例信息
                self.keyboard_ratios = {
                    'key_width_ratio': avg_width / img_width,
                    'key_height_ratio': avg_height / img_height,
                    'row_spacing_ratio': row_spacing / img_height,
                    'col_spacing_ratio': col_spacing / img_width,
                    'image_width': img_width,
                    'image_height': img_height
                }
                
    def _group_keys_by_rows(self, keys: List[Dict]) -> List[List[Dict]]:
        """按行分组按键"""
        if not keys:
            return []
            
        # 按y坐标排序
        sorted_keys = sorted(keys, key=lambda k: k['y'])
        
        rows = []
        current_row = [sorted_keys[0]]
        row_threshold = sorted_keys[0]['height'] * 0.5  # 行间距阈值
        
        for key in sorted_keys[1:]:
            if key['y'] - current_row[-1]['y'] <= row_threshold:
                current_row.append(key)
            else:
                rows.append(sorted(current_row, key=lambda k: k['x']))
                current_row = [key]
                
        if current_row:
            rows.append(sorted(current_row, key=lambda k: k['x']))
            
        return rows
        
    def _calculate_row_spacing(self, rows: List[List[Dict]]) -> float:
        """计算行间距"""
        if len(rows) < 2:
            return 0
            
        spacings = []
        for i in range(len(rows) - 1):
            row1_center_y = np.mean([key['y'] for key in rows[i]])
            row2_center_y = np.mean([key['y'] for key in rows[i+1]])
            spacing = row2_center_y - row1_center_y
            spacings.append(spacing)
            
        return np.mean(spacings) if spacings else 0
        
    def _calculate_col_spacing(self, rows: List[List[Dict]]) -> float:
        """计算列间距"""
        spacings = []
        for row in rows:
            if len(row) > 1:
                for i in range(len(row) - 1):
                    spacing = row[i+1]['x'] - row[i]['x']
                    spacings.append(spacing)
                    
        return np.mean(spacings) if spacings else 0
        
    def compare_keyboards(self, reference_keys: List[Dict], current_keys: List[Dict], 
                         current_image_shape: Tuple[int, int, int]) -> List[KeyIssue]:
        """
        对比两个键盘的按键配置 - 基于比例
        """
        issues = []
        
        # 检查按键数量
        if len(current_keys) < len(reference_keys):
            issues.append(KeyIssue(
                issue_type=KeyIssueType.MISSING,
                position=(0, 0),
                description=f"检测到 {len(reference_keys) - len(current_keys)} 个按键缺失",
                confidence=0.9
            ))
            
        # 为每个参考按键找到最近的当前按键
        used_current_indices = set()
        
        for ref_key in reference_keys:
            best_match_idx = None
            best_match_distance = float('inf')
            best_match_score = 0
            
            for i, curr_key in enumerate(current_keys):
                if i in used_current_indices:
                    continue
                    
                # 计算位置距离（基于比例）
                distance = self._calculate_ratio_distance(ref_key, curr_key, current_image_shape)
                
                # 计算相似度分数
                similarity = self._calculate_key_similarity_ratio(ref_key, curr_key, current_image_shape)
                
                # 综合评分
                if distance < self.position_tolerance_ratio and similarity > self.match_threshold:
                    if distance < best_match_distance:
                        best_match_idx = i
                        best_match_distance = distance
                        best_match_score = similarity
                        
            if best_match_idx is not None:
                used_current_indices.add(best_match_idx)
                curr_key = current_keys[best_match_idx]
                
                # 检查具体问题
                issues.extend(self._check_key_issues_ratio(ref_key, curr_key, current_image_shape))
            else:
                # 未找到匹配的按键
                issues.append(KeyIssue(
                    issue_type=KeyIssueType.MISSING,
                    position=(ref_key['x'], ref_key['y']),
                    description=f"在位置 ({ref_key['x']}, {ref_key['y']}) 处缺失按键",
                    confidence=0.8
                ))
                
        return issues
        
    def _calculate_ratio_distance(self, key1: Dict, key2: Dict, image_shape: Tuple[int, int, int]) -> float:
        """计算基于比例的距离"""
        img_height, img_width = image_shape[:2]
        
        # 计算相对位置
        x1_ratio = key1['x'] / img_width
        y1_ratio = key1['y'] / img_height
        x2_ratio = key2['x'] / img_width
        y2_ratio = key2['y'] / img_height
        
        # 计算比例距离
        distance = np.sqrt((x1_ratio - x2_ratio)**2 + (y1_ratio - y2_ratio)**2)
        
        return distance
        
    def _calculate_key_similarity_ratio(self, key1: Dict, key2: Dict, image_shape: Tuple[int, int, int]) -> float:
        """计算基于比例的按键相似度"""
        img_height, img_width = image_shape[:2]
        
        # 位置相似度（基于比例）
        pos_distance = self._calculate_ratio_distance(key1, key2, image_shape)
        pos_similarity = max(0, 1 - pos_distance / self.position_tolerance_ratio)
        
        # 大小相似度（基于比例）
        key1_width_ratio = key1['width'] / img_width
        key1_height_ratio = key1['height'] / img_height
        key2_width_ratio = key2['width'] / img_width
        key2_height_ratio = key2['height'] / img_height
        
        width_diff = abs(key1_width_ratio - key2_width_ratio) / key1_width_ratio
        height_diff = abs(key1_height_ratio - key2_height_ratio) / key1_height_ratio
        size_similarity = max(0, 1 - (width_diff + height_diff) / 2)
        
        # 角度相似度
        angle_diff = abs(key1['angle'] - key2['angle']) % 360
        if angle_diff > 180:
            angle_diff = 360 - angle_diff
        angle_similarity = max(0, 1 - angle_diff / 90)
        
        # 综合相似度
        total_similarity = (pos_similarity * 0.4 + size_similarity * 0.3 + angle_similarity * 0.3)
        
        return total_similarity
        
    def _check_key_issues_ratio(self, ref_key: Dict, curr_key: Dict, image_shape: Tuple[int, int, int]) -> List[KeyIssue]:
        """检查按键的具体问题 - 基于比例"""
        issues = []
        img_height, img_width = image_shape[:2]
        
        # 检查位置偏差（基于比例）
        pos_distance = self._calculate_ratio_distance(ref_key, curr_key, image_shape)
        
        if pos_distance > self.position_tolerance_ratio:
            issues.append(KeyIssue(
                issue_type=KeyIssueType.MISPLACED,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键位置偏差 {pos_distance:.3f} (比例)",
                confidence=min(0.9, pos_distance / self.position_tolerance_ratio)
            ))
            
        # 检查旋转
        angle_diff = abs(ref_key['angle'] - curr_key['angle']) % 360
        if angle_diff > 180:
            angle_diff = 360 - angle_diff
            
        if angle_diff > self.angle_tolerance:
            issues.append(KeyIssue(
                issue_type=KeyIssueType.ROTATED,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键旋转 {angle_diff:.1f} 度",
                confidence=min(0.9, angle_diff / 45)
            ))
            
        # 检查大小差异（基于比例）
        ref_width_ratio = ref_key['width'] / img_width
        ref_height_ratio = ref_key['height'] / img_height
        curr_width_ratio = curr_key['width'] / img_width
        curr_height_ratio = curr_key['height'] / img_height
        
        width_diff = abs(ref_width_ratio - curr_width_ratio) / ref_width_ratio
        height_diff = abs(ref_height_ratio - curr_height_ratio) / ref_height_ratio
        
        if width_diff > self.size_tolerance_ratio or height_diff > self.size_tolerance_ratio:
            issues.append(KeyIssue(
                issue_type=KeyIssueType.WRONG_KEY,
                position=(curr_key['x'], curr_key['y']),
                description=f"按键大小异常 (宽差: {width_diff:.1%}, 高差: {height_diff:.1%})",
                confidence=max(width_diff, height_diff)
            ))
            
        return issues
        
    def visualize_detection(self, image: npt.NDArray, keys: List[Dict], issues: List[KeyIssue]) -> npt.NDArray:
        """可视化检测结果 - 增强版本"""
        vis_image = image.copy()
        
        # 绘制检测到的按键
        for i, key in enumerate(keys):
            # 绘制按键边界框
            cv2.rectangle(
                vis_image,
                (key['x'] - key['width']//2, key['y'] - key['height']//2),
                (key['x'] + key['width']//2, key['y'] + key['height']//2),
                (0, 255, 0), 2
            )
            
            # 绘制按键中心点
            cv2.circle(vis_image, (key['x'], key['y']), 3, (0, 0, 255), -1)
            
            # 绘制按键编号（可选）
            cv2.putText(
                vis_image, str(i+1), 
                (key['x'] - 5, key['y'] - 5),
                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1
            )
            
        # 绘制问题标记
        issue_count = 0
        for issue in issues:
            x, y = issue.position
            
            if issue.issue_type == KeyIssueType.MISSING:
                color = (0, 0, 255)  # 红色 - 缺失
                symbol = "X"
                thickness = 3
            elif issue.issue_type == KeyIssueType.MISPLACED:
                color = (0, 255, 255)  # 黄色 - 位置错误
                symbol = "M"
                thickness = 2
            elif issue.issue_type == KeyIssueType.ROTATED:
                color = (255, 255, 0)  # 青色 - 旋转
                symbol = "R"
                thickness = 2
            else:
                color = (255, 0, 255)  # 紫色 - 错误按键
                symbol = "W"
                thickness = 2
                
            # 绘制问题标记圆圈
            cv2.circle(vis_image, (x, y), 20, color, thickness)
            
            # 绘制问题符号
            cv2.putText(
                vis_image, symbol, (x - 8, y + 8),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2
            )
            
            # 绘制问题编号
            issue_count += 1
            cv2.putText(
                vis_image, str(issue_count), 
                (x + 15, y - 15),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1
            )
            
        # 添加统计信息
        total_keys = len(keys)
        total_issues = len(issues)
        
        # 统计不同类型的问题
        missing_count = len([i for i in issues if i.issue_type == KeyIssueType.MISSING])
        misplaced_count = len([i for i in issues if i.issue_type == KeyIssueType.MISPLACED])
        rotated_count = len([i for i in issues if i.issue_type == KeyIssueType.ROTATED])
        wrong_count = len([i for i in issues if i.issue_type == KeyIssueType.WRONG_KEY])
        
        # 绘制信息面板背景
        info_height = 120
        cv2.rectangle(vis_image, (10, 10), (400, info_height), (0, 0, 0), -1)
        cv2.rectangle(vis_image, (10, 10), (400, info_height), (255, 255, 255), 2)
        
        # 绘制统计信息
        y_offset = 30
        cv2.putText(vis_image, f"检测到按键: {total_keys}", (20, y_offset), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        y_offset += 25
        
        cv2.putText(vis_image, f"总问题数: {total_issues}", (20, y_offset), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        y_offset += 25
        
        if missing_count > 0:
            cv2.putText(vis_image, f"缺失按键: {missing_count}", (20, y_offset), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
            y_offset += 20
        
        if misplaced_count > 0:
            cv2.putText(vis_image, f"位置错误: {misplaced_count}", (20, y_offset), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)
            y_offset += 20
        
        if rotated_count > 0:
            cv2.putText(vis_image, f"方向错误: {rotated_count}", (20, y_offset), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 2)
            y_offset += 20
        
        if wrong_count > 0:
            cv2.putText(vis_image, f"错误按键: {wrong_count}", (20, y_offset), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 255), 2)
        
        # 绘制图例
        legend_x = image.shape[1] - 200
        legend_y = 30
        cv2.rectangle(vis_image, (legend_x, 10), (image.shape[1] - 10, 150), (0, 0, 0), -1)
        cv2.rectangle(vis_image, (legend_x, 10), (image.shape[1] - 10, 150), (255, 255, 255), 2)
        
        cv2.putText(vis_image, "图例:", (legend_x + 10, legend_y), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        legend_y += 25
        
        cv2.putText(vis_image, "X - 缺失按键", (legend_x + 10, legend_y), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
        legend_y += 20
        
        cv2.putText(vis_image, "M - 位置错误", (legend_x + 10, legend_y), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)
        legend_y += 20
        
        cv2.putText(vis_image, "R - 方向错误", (legend_x + 10, legend_y), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 2)
        legend_y += 20
        
        cv2.putText(vis_image, "W - 错误按键", (legend_x + 10, legend_y), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 255), 2)
        
        return vis_image
